package com.km.algorithm

import org.apache.commons.math3.distribution.ChiSquaredDistribution
import org.apache.spark.mllib.linalg.{Vector, Vectors}

/**
  * Created by lenovo on 2017/4/18.
  */
object ComputeUtil {

  private def methodFromString(methodName: String): Method = {
    methodName match {
      case PEARSON.name => PEARSON
      case _ => throw new IllegalArgumentException("Unrecognized method for Chi squared test.")
    }
  }

  val PEARSON = new Method("pearson", (observed: Double, expected: Double) => {
    val dev = observed - expected
    dev * dev / expected
  })

  case class Method(name: String, chiSqFunc: (Double, Double) => Double)

  def getPValue(observed: Vector,
                expected: Vector): Double = {
    val obsArr = observed.toArray
    //    val expArr = Array(1.2,1.9,1.0,0.3,1.6,1.4,0.9,1.1,1.1,0.5);
    //    val expected: Vector = Vectors.dense(1.2,1.9,1.0,0.3,1.6,1.4,0.9,1.1,1.1,0.5)
    val methodName: String = "pearson"
    //    val method = PEARSON
    val method = methodFromString(methodName)
    val size = obsArr.size


    val expArr = if (expected.size == 0) Array.tabulate(size)(_ => 1.0 / size) else expected.toArray
    val obsSum = obsArr.sum
    val expSum = if (expected.size == 0.0) 1.0 else expArr.sum
    val scale = if (math.abs(obsSum - expSum) < 1e-7) 1.0 else obsSum / expSum
    val statistic = obsArr.zip(expArr).foldLeft(0.0) { case (stat, (obs, exp)) =>
      //      if (exp == 0.0) {
      //        if (obs == 0.0) {
      //          throw new IllegalArgumentException("Chi-squared statistic undefined for input vectors due"
      //            + " to 0.0 values in both observed and expected.")
      //        } else {
      //          return new ChiSqTestResult(0.0, size - 1, Double.PositiveInfinity, PEARSON.name,
      //            NullHypothesis.goodnessOfFit.toString)
      //        }
      //      }
      if (scale == 1.0) {
        stat + method.chiSqFunc(obs, exp)
      } else {
        stat + method.chiSqFunc(obs, exp * scale)
      }
    }
    val df2 = obsArr.length - 1
    val pValue = 1.0 - new ChiSquaredDistribution(df2).cumulativeProbability(statistic)
    println("// pValue")
    pValue

  }

  def main(args: Array[String]): Unit = {
    val v1: Vector = Vectors.dense(1.3, 1.5, 1.5, 0.7, 1.1, 1.0, 0.8, 1.9, 1.7, 0.2)

    val v2 = Vectors.dense(1.2, 1.9, 1.0, 0.3, 1.6, 1.4, 0.9, 1.1, 1.1, 0.5)

    println(getPValue(v1, v2));

    val v3 = Vectors.dense(1, 2, 3, 4, 5)

    val v4 = Vectors.dense(3, 5, 9, 3, 2)

    println(getPValue(v3, v4));
  }


}
